# TODO 需要命令行执行
import numpy as np
import onnx
from PIL import Image
from tvm import relay

onnx_model = onnx.load('inception-v4.onnx')
img = Image.open('cat.jpg').resize((299, 299))

# 以下的图片读取仅仅是为了测试
img = np.array(img).transpose((2, 0, 1)).astype('float32')
img = img/255.0    # remember pytorch tensor is 0-1
x = img[np.newaxis, :]

# 这里首先在PC的CPU上进行测试 所以使用LLVM进行导出
# target = 'llvm'

# 在PC的GPU上进行测试
# gpu
target = 'cuda'
input_name = 'input'
shape_dict = {input_name: x.shape}
sym, params = relay.frontend.from_onnx(onnx_model, shape_dict)


# 这里利用TVM构建出优化后模型的信息
with relay.build_config(opt_level=3):
    graph, lib, params = relay.build_module.build(mod=sym,
                                                  target=target,
                                                  params=params)

# 下面的函数导出我们需要的动态链接库 地址可以自己定义
print("Output model files")
libpath = "inception_v4.so"
lib.export_library(libpath)

# 下面的函数导出我们神经网络的结构，使用json文件保存
graph_json_path = "inception_v4.json"
with open(graph_json_path, 'w') as fo:
    fo.write(graph)

# 下面的函数中我们导出神经网络模型的权重参数
param_path = "inception_v4.params"
with open(param_path, 'wb') as fo:
    fo.write(relay.save_param_dict(params))
